package cn.org.intelli.zjgflink.util;


import cn.org.intelli.zjgflink.service.HdfsOpService;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.client.deployment.ClusterSpecification;
import org.apache.flink.client.deployment.application.ApplicationConfiguration;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.ClusterClientProvider;
import org.apache.flink.configuration.*;
import org.apache.flink.runtime.jobgraph.SavepointConfigOptions;
import org.apache.flink.yarn.YarnClientYarnClusterInformationRetriever;
import org.apache.flink.yarn.YarnClusterDescriptor;
import org.apache.flink.yarn.configuration.YarnConfigOptions;
import org.apache.flink.yarn.configuration.YarnDeploymentTarget;
import org.apache.flink.yarn.configuration.YarnLogConfigUtil;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.Collections;

@Component
public class FlinkJobDeploy {
    private static final Logger log = LoggerFactory.getLogger(cn.org.intelli.zjgflink.util.FlinkJobDeploy.class);

    @Value("${flink.hdfs.jarDir}")
    private String flinkHdfsJarDir;

    @Value("${flink.hdfs.libs}")
    private String flinkHdfsLibs;

    @Value("${flink.hdfs.distJar}")
    private String flinkHdfsDistJar;

    @Value("${flink.hdfs.job}")
    private String flinkHdfsJob;

    @Value("${flink.conf.home}")
    private String flinkConfHome;

    @Autowired
    private HdfsOpService hdfsOpService;

    public String deployFlinkJob(String flinkJobId, String flinkJobName,
                                 Integer slotNumber, String mainFunction, String jarFileName,
                                 String savePointDir, Integer parallelism,
                                 String jobManagerProcessSize, String taskManagerProcessSize,
                                 String[] args) throws Exception {
        System.setProperty("HADOOP_USER_NAME", "root");
        log.info("@FlinkJobDeploy: deploy flink job: flinkJobId is {}," +
                        " flinkJobName is {},slotNumber is {},mainFunction is {},jarFileName is {}," +
                        "savePointDir is {},args are {}",
                flinkJobId, flinkJobName, slotNumber, mainFunction,
                jarFileName, savePointDir, args);
        YarnClient yarnClient = YarnClient.createYarnClient();
        YarnConfiguration yarnConfiguration = new YarnConfiguration();
        yarnClient.init(yarnConfiguration);
        yarnClient.start();
        try {
            YarnClientYarnClusterInformationRetriever yarnClientYarnClusterInformationRetriever =
                    YarnClientYarnClusterInformationRetriever.create(yarnClient);
            Configuration flinkConfiguration = GlobalConfiguration.loadConfiguration(this.flinkConfHome);
            flinkConfiguration.set(CheckpointingOptions.CHECKPOINTS_DIRECTORY, this.flinkHdfsJob + "/" + flinkJobId);

            if (StringUtils.isNotEmpty(savePointDir)) {
                flinkConfiguration.set(SavepointConfigOptions.SAVEPOINT_PATH, savePointDir);
            } else {
                String filePath = this.flinkHdfsJob + "/" + flinkJobId;
                Boolean isFilePathExist = this.hdfsOpService.checkFolderExisted(filePath);
                if (isFilePathExist.booleanValue()) {
                    FileStatus[] fileStatuses = this.hdfsOpService.getFileStatusInFolder(filePath);
                    if (fileStatuses != null && fileStatuses.length != 0) {
                        Long minTime = Long.valueOf(0L);
                        String lastChkPointFolderName = null;
                        for (FileStatus folderStatus : fileStatuses) {
                            String folderName = folderStatus.getPath().getName();
                            if (!StringUtils.startsWith(folderName, "savepoint-")) {
                                Long lastModified = Long.valueOf(folderStatus.getModificationTime());
                                if (lastModified.longValue() > minTime.longValue()) {
                                    minTime = lastModified;
                                    lastChkPointFolderName = folderStatus.getPath().getName();
                                }
                            }
                        }
                        if (StringUtils.isNotEmpty(lastChkPointFolderName)) {
                            String ckPath = filePath + "/" + lastChkPointFolderName;
                            minTime = Long.valueOf(0L);
                            String ckFolderName = null;
                            FileStatus[] tmpFileStatuses = this.hdfsOpService.getFileStatusInFolder(ckPath);
                            for (FileStatus fileStatus : tmpFileStatuses) {
                                if (StringUtils.startsWith(fileStatus.getPath().getName(), "chk-"))
                                    if (fileStatus.getModificationTime() > minTime.longValue()) {
                                        minTime = Long.valueOf(fileStatus.getModificationTime());
                                        ckFolderName = fileStatus.getPath().getName();
                                    }
                            }
                            if (StringUtils.isNotEmpty(ckFolderName)) {
                                String finalCkPath = ckPath + "/" + ckFolderName;
                                flinkConfiguration.set(SavepointConfigOptions.SAVEPOINT_PATH, finalCkPath);
                            }
                        }
                    }
                } else {
                    log.info("HDFS{}", filePath);
                }
            }

            flinkConfiguration.set(PipelineOptions.JARS, Collections.singletonList(this.flinkHdfsJarDir + "/" + jarFileName));
            Path remoteLib = new Path(this.flinkHdfsLibs);
            flinkConfiguration.set(YarnConfigOptions.PROVIDED_LIB_DIRS, Collections.singletonList(remoteLib.toString()));
            flinkConfiguration.set(YarnConfigOptions.FLINK_DIST_JAR, this.flinkHdfsDistJar);
            flinkConfiguration.set(DeploymentOptions.TARGET, YarnDeploymentTarget.APPLICATION.getName());

            if (parallelism == null || parallelism.intValue() == 0)
                parallelism = Integer.valueOf(1);
            flinkConfiguration.set(CoreOptions.DEFAULT_PARALLELISM, parallelism);
            flinkConfiguration.set(TaskManagerOptions.NUM_TASK_SLOTS, slotNumber);
            if (StringUtils.isNotEmpty(taskManagerProcessSize))
                flinkConfiguration.set(TaskManagerOptions.TOTAL_PROCESS_MEMORY, MemorySize.parse(taskManagerProcessSize));
            if (StringUtils.isNotEmpty(jobManagerProcessSize))
                flinkConfiguration.set(JobManagerOptions.TOTAL_PROCESS_MEMORY, MemorySize.parse(jobManagerProcessSize));
            flinkConfiguration.set(YarnConfigOptions.APPLICATION_NAME, flinkJobName);

            ClusterSpecification clusterSpecification = (new ClusterSpecification.ClusterSpecificationBuilder()).createClusterSpecification();
            ApplicationConfiguration appConfig = new ApplicationConfiguration(args, mainFunction);
            YarnLogConfigUtil.setLogConfigFileInConfig(flinkConfiguration, this.flinkConfHome);
            YarnClusterDescriptor yarnClusterDescriptor = new YarnClusterDescriptor(flinkConfiguration, yarnConfiguration,
                    yarnClient, yarnClientYarnClusterInformationRetriever, true);

            ClusterClientProvider<ApplicationId> clusterClientProvider = yarnClusterDescriptor.deployApplicationCluster(clusterSpecification, appConfig);
            ClusterClient<ApplicationId> clusterClient = clusterClientProvider.getClusterClient();
            ApplicationId applicationId = clusterClient.getClusterId();
            log.info("@FlinkJobDeploy: deploy flink job {} successful, applicationId is {}", flinkJobName, applicationId);
            return applicationId.toString();
        } catch (Exception e) {
            e.printStackTrace();
            throw new Exception(e);
        } finally {
            yarnClient.stop();
        }
    }
}
